#!/usr/bin/env python3
"""
PDF 转 Markdown 脚本
使用 PyMuPDF (fitz) 提取 PDF 内容并转换为 Markdown 格式
"""

import sys
import argparse
from pathlib import Path

try:
    import fitz  # PyMuPDF
except ImportError:
    print("错误: 需要安装 PyMuPDF 库")
    print("请运行: pip install PyMuPDF")
    sys.exit(1)


def extract_text_with_formatting(page):
    """从 PDF 页面提取文本，尝试保留格式"""
    blocks = page.get_text("dict")["blocks"]
    markdown_lines = []

    for block in blocks:
        if block["type"] == 0:  # 文本块
            for line in block["lines"]:
                line_text = ""
                for span in line["spans"]:
                    text = span["text"].strip()
                    if not text:
                        continue

                    # 根据字体大小判断是否为标题
                    font_size = span["size"]
                    if font_size > 16:
                        line_text += f"# {text}"
                    elif font_size > 14:
                        line_text += f"## {text}"
                    elif font_size > 12:
                        line_text += f"### {text}"
                    else:
                        line_text += text
                    line_text += " "

                if line_text.strip():
                    markdown_lines.append(line_text.strip())

        elif block["type"] == 1:  # 图片块
            markdown_lines.append("\n[图片]\n")

    return "\n\n".join(markdown_lines)


def pdf_to_markdown(pdf_path: Path, output_path: Path = None, simple_mode: bool = False):
    """
    将 PDF 转换为 Markdown

    Args:
        pdf_path: PDF 文件路径
        output_path: 输出 Markdown 文件路径（可选）
        simple_mode: 简单模式，只提取纯文本
    """
    if not pdf_path.exists():
        print(f"错误: PDF 文件不存在: {pdf_path}")
        sys.exit(1)

    # 确定输出路径
    if output_path is None:
        output_path = pdf_path.with_suffix('.md')

    print(f"正在转换: {pdf_path}")
    print(f"输出到: {output_path}")

    try:
        # 打开 PDF
        doc = fitz.open(pdf_path)
        markdown_content = []

        # 添加文档信息
        markdown_content.append(f"# {pdf_path.stem}\n")
        markdown_content.append(f"转换自: {pdf_path.name}\n")
        markdown_content.append(f"总页数: {len(doc)}\n")
        markdown_content.append("---\n")

        # 处理每一页
        for page_num in range(len(doc)):
            page = doc[page_num]
            print(f"处理第 {page_num + 1}/{len(doc)} 页...")

            markdown_content.append(f"\n## 第 {page_num + 1} 页\n")

            if simple_mode:
                # 简单模式：只提取纯文本
                text = page.get_text()
                markdown_content.append(text)
            else:
                # 格式化模式：尝试保留格式
                formatted_text = extract_text_with_formatting(page)
                markdown_content.append(formatted_text)

            markdown_content.append("\n---\n")

        # 写入文件
        output_path.write_text("\n".join(markdown_content), encoding='utf-8')

        print(f"✓ 转换完成!")
        print(f"输出文件: {output_path}")
        print(f"文件大小: {output_path.stat().st_size / 1024:.2f} KB")

        doc.close()

    except Exception as e:
        print(f"错误: 转换失败 - {e}")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description='将 PDF 文件转换为 Markdown 格式',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  python pdf_to_markdown.py input.pdf
  python pdf_to_markdown.py input.pdf -o output.md
  python pdf_to_markdown.py input.pdf --simple
        '''
    )

    parser.add_argument(
        'pdf_file',
        type=Path,
        help='输入的 PDF 文件路径'
    )

    parser.add_argument(
        '-o', '--output',
        type=Path,
        help='输出的 Markdown 文件路径（默认与 PDF 同名）'
    )

    parser.add_argument(
        '--simple',
        action='store_true',
        help='简单模式：只提取纯文本，不保留格式'
    )

    args = parser.parse_args()

    pdf_to_markdown(args.pdf_file, args.output, args.simple)


if __name__ == '__main__':
    main()
